/**
    Copyright (c) 2010 yakiimo02
    Distributed under the New BSD License.
    See included license.txt or http://www.yakiimo3d.com/NewBSDLicense.txt
**/

#include "OITShaderCommon.h"

StructuredBuffer<FragmentLink> FragmentLinkSRV            : register(t0);
Buffer<uint> StartOffsetSRV                                : register(t1);

struct QuadVSinput
{
    float4 pos : POSITION;
};

struct QuadVS_Output
{
    float4 pos : SV_POSITION;              
};

/** 
    Draw full screen quad.
*/
QuadVS_Output QuadVS( QuadVSinput Input )
{
    QuadVS_Output Output;
    Output.pos = Input.pos;
    return Output;
}

// Max hardcoded.
// We're in trouble if the fragment linked list is larger than 32...
#define TEMPORARY_BUFFER_MAX        32

/**
    Sort and render fragments.
*/
float4 SortFragmentsPS( QuadVS_Output input ) : SV_Target0
{    
    // index to current pixel.
    uint nIndex = (uint)input.pos.y * g_nFrameWidth + (uint)input.pos.x;
    
    FragmentData aData[ TEMPORARY_BUFFER_MAX ];            // temporary buffer
    int anIndex[ TEMPORARY_BUFFER_MAX ];                // index array for the tempory buffer
    uint nNumFragment = 0;                                // number of fragments in current pixel's linked list.
    uint nNext = StartOffsetSRV[ nIndex ];                // get first fragment from the start offset buffer.
    
    // early exit if no fragments in the linked list.
    if( nNext == 0xFFFFFFFF ) {
        return float4( 0.0, 0.0, 0.0, 0.0 );
    }
    
    // Read and store linked list data to the temporary buffer.
    while( nNext != 0xFFFFFFFF ) 
    {
        FragmentLink element = FragmentLinkSRV[nNext];
        aData[ nNumFragment ] = element.fragmentData;
        anIndex[ nNumFragment ] = nNumFragment;
        ++nNumFragment;
        nNext = element.nNext;
    }

    uint N2 = 1 << (int)(ceil(log2(nNumFragment)));
    
    // bitonic sort implementation needs on pow2 data.
    for(int i = nNumFragment; i < N2; i++)
    {
        anIndex[i] = i;
        aData[i].fDepth = 1.1f;
    }

    // Unoptimized sorting. (Bitonic Sort)
    
    // loop from Merge( 2 ) to Merge( nCount )
    for( int nMergeSize=2; nMergeSize<=N2; nMergeSize=nMergeSize*2 )
    {
        // Merge( nCount ) requires log2( nCount ) merges. Merge( nCount/2 ) -> Merge( 2 )
        for( int nMergeSubSize=nMergeSize>>1; nMergeSubSize>0; nMergeSubSize=nMergeSubSize>>1 )
        {
            // compare and swap elements
            for( int nElem=0; nElem<N2; ++nElem )
            {
                int nSwapElem = nElem^nMergeSubSize;
                // check to make sure to only swap once
                if( nSwapElem > nElem )
                {
                    // sort in increasing order
                    if ( ( nElem & nMergeSize ) == 0 && aData[ anIndex[ nElem ] ].fDepth>aData[ anIndex[ nSwapElem ] ].fDepth ) 
                    {
                        int temp = anIndex[ nElem ];
                        anIndex[ nElem ] = anIndex[ nSwapElem ];
                        anIndex[ nSwapElem ] = temp;
                    }

                    // sort in descending order
                    if ( ( nElem & nMergeSize ) != 0 && aData[ anIndex[ nElem ] ].fDepth<aData[ anIndex[ nSwapElem ] ].fDepth ) 
                    {
                        int temp = anIndex[ nElem ];
                        anIndex[ nElem ] = anIndex[ nSwapElem ];
                        anIndex[ nSwapElem ] = temp;
                    }
                }
            }
        }
    }

    // Now blend the sorted fragments.
    float4 result = 0.0f;
    for( int x = nNumFragment-1; x >= 0; --x )
    {
        uint nColor = aData[ anIndex[ x ] ].nColor;
        float4 color;        
        color.r = ( (nColor >> 0) & 0xFF ) / 255.0f;
        color.g = ( (nColor >> 8) & 0xFF ) / 255.0f;
        color.b = ( (nColor >> 16) & 0xFF ) / 255.0f;
        color.a = ( (nColor >> 24) & 0xFF ) / 255.0f;
        result = lerp( result, color, color.a );
    }
    
    result.a = 1.0f;
    return result;
}
